package com.atguigu.bigdata.spark.sql;

import org.apache.spark.SparkConf;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.expressions.MutableAggregationBuffer;
import org.apache.spark.sql.expressions.UserDefinedAggregateFunction;
import org.apache.spark.sql.types.DataType;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import org.sparkproject.guava.collect.Lists;

import java.util.ArrayList;
import java.util.List;

/*
    自定义聚合函数类：计算年龄的平均值
 */

public class Spark03_SparkSQL_UDAF_JAVA {
    public static void main(String[] args) {
        SparkConf conf = new  SparkConf().setMaster("local[*]").setAppName("sparkSQL");
        SparkSession spark = SparkSession.builder().config(conf).getOrCreate();
        Dataset<Row> df   = spark.read().json("datas/user.json");

        df.createOrReplaceTempView("user");
        spark.udf().register("ageAvg", new MyAvgUDAF());
        spark.sql("select ageAvg(age) from user").show();

        spark.close();
    }


}

class  MyAvgUDAF extends UserDefinedAggregateFunction {
    // 输入数据的结构
    @Override
    public StructType inputSchema() {
        List<StructField> fields = new ArrayList<>();
        fields.add(DataTypes.createStructField("age", DataTypes.LongType,true));
        return DataTypes.createStructType(fields);
    }
    // 缓冲区数据的结构 : Buffer
    @Override
    public StructType bufferSchema() {
        List<StructField> fields = new ArrayList<>();
        fields.add(DataTypes.createStructField("count", DataTypes.LongType,true));
        fields.add(DataTypes.createStructField("total", DataTypes.LongType,true));
        return DataTypes.createStructType(fields);
    }
    // 函数计算结果的数据类型：Out
    @Override
    public DataType dataType() {
        return DataTypes.LongType;
    }
    // 函数的稳定性
    @Override
    public boolean deterministic() {
        return true;
    }
    // 缓冲区初始化
    @Override
    public void initialize(MutableAggregationBuffer buffer) {
        //0 和 1代表位置 对应 缓冲区数据结构的list里面的内容
        buffer.update(0,0L);
        buffer.update(1,0L);
    }
    // 根据输入的值更新缓冲区数据
    @Override
    public void update(MutableAggregationBuffer buffer, Row input) {
        buffer.update(0, buffer.getLong(0)+input.getLong(0));
        buffer.update(1, buffer.getLong(1)+1);
    }
    // 缓冲区数据合并
    @Override
    public void merge(MutableAggregationBuffer buffer1, Row buffer2) {
        buffer1.update(0, buffer1.getLong(0) + buffer2.getLong(0));
        buffer1.update(1, buffer1.getLong(1) + buffer2.getLong(1));
    }
    //计算并返回最终结果
    @Override
    public Object evaluate(Row buffer) {
        return buffer.getLong(0)/buffer.getLong(1);
    }
}
